import os
import json

import pandas as pd


def read_task_from_file(file_path="./task/example.json"):
    with open(file_path) as file:
        task = json.load(file)

    return task


def get_config_value(config, key):
    dict_keys = ["reporters", "hs_code", "partners", "partner2nds", "hs_code2industry", "region2id", "periods"]
    if key not in dict_keys:
        return config[key]
    
    if key == "periods":
        try:
            if type(config[key]) == list:
                return config[key]
            
            elif type(config[key]) == str:
                if config[key][4] == "-":
                    start_year = config[key][:4]
                    end_year = config[key][5:9]

                    return list(range(int(start_year), int(end_year) + 1))
        except:
            print("Can't read periods, please check.")
            print("Please write periods with a list like [2012, 2015], or a string like '2012-2015'.")
            exit(4)

    if config[key]["src"] == "csv":
        csv_path = config[key]["data"]
        if key == "hs_code":
            return read_hs_code_from_file(csv_path)
        
        if key == "hs_code2industry":
            return read_hs_code2industry_from_file(csv_path)
        
        if key == "region2id":
            return read_region2id_from_file(csv_path)

    elif config[key]["src"] == "value":
        return config[key]["data"]
    
    elif config[key]["src"] == "industry":
        csv_path = config[key]["data"]
        return read_partners_by_industry_from_file(csv_path)
    
    else:
        raise ValueError(f"Unsupport data source type: {config[key]['src']}.")


def read_hs_code_from_file(file_path="./data/hs_code.csv"):
    df = pd.read_csv(file_path, header=None, dtype=str) 
    # read hs_code as a string
    hs_codes = df[0].to_list()
    
    return hs_codes


def read_hs_code2industry_from_file(file_path="./data/hs_code2industry.csv"):
    # read hs_code as a string
    df = pd.read_csv(file_path, dtype=str) 

    # one hs_code may have multiple industries type
    hs_code2industry = df.groupby('hs_code')['industry'].apply(list).to_dict() 

    return hs_code2industry


def read_partners_by_industry_from_file(file_path="./data/industry2partners.csv"):
    df = pd.read_csv(file_path, header=0, dtype=str)

    industry2partners = {}
    
    for col in df.columns:  
        col_list = df[col].dropna().tolist()  
        industry2partners[col] = col_list

    return industry2partners


def read_region2id_from_file(file_path="./data/region2id.csv"):
    df = pd.read_csv(file_path, dtype=str) 

    short_df = df.set_index('region') 
    short_region2id = short_df['id'].to_dict()
    chinese_df = df.set_index('chinese_name')
    chinese_region2id = chinese_df['id'].to_dict()

    region2id = {
        "short": short_region2id,
        "chinese": chinese_region2id
    }

    return region2id


def read_auth_from_file(file_path="./.auth"):
    with open(file_path) as file:
        auth = file.readline().strip()
    
    return auth


def get_region_id(region_name, region_ids):
    if region_name in region_ids["short"]:
        return region_ids["short"][region_name]
    
    if region_name in region_ids["chinese"]:
        return region_ids["chinese"][region_name]

    error_str = f"The region {region_name} id is unknown, please add the id to the code."
    raise ValueError(error_str)


def save_no_industry_codes2txt(no_industry_code, save_path="./no_ind_code.txt"):
    with open(save_path, "w") as file:
        for code in no_industry_code:
            file.write(str(code) + "\n")


def quick_save(save_info, save_path="./quick_save.txt"):
    with open(save_path, "a") as file:
        file.write(str(save_info) + "\n")


def quick_load(save_path="./quick_save.txt"):
    if not os.path.exists(save_path):
        return []

    with open(save_path, "r") as file:
        lines = file.readlines()

    # reporters # periods # hs_codes # partners # partner2nds

    saved_infos = [line.strip() for line in lines]

    return saved_infos


def save_raw_data2csv(data, csv_path):
    if data is None:
        print("Data is empty.")
    
    df = pd.DataFrame(data)
    df.to_csv(csv_path, index=False, encoding="utf-8")
    print("Data saved to " + csv_path + ".")


def generate_output_filename(format, hs_code="", reporters=[], periods="", industry=""):
    filename = format

    reporters_str = str(reporters)
    if type(reporters) == list:
        reporters_str = '-'.join(reporters)

    filename = filename.replace(r"%reporters%", reporters_str)
    filename = filename.replace(r"%hs_code%", hs_code)
    filename = filename.replace(r"%periods%", str(periods))
    filename = filename.replace(r"%industry%", str(industry))

    return filename


if __name__ == "__main__":
    # test
    pass
